from data_structure.stack import Stack


class Bi_tree(object):
    def __init__(self, data=5, lchhid=None, rchid=None):
        self.data = data
        self.lchild = lchhid
        self.rchild = rchid

    def pre_order_recursion(self, root):
        if root is None:
            return
        print(root.data, end=' ')
        self.pre_order_recursion(root.lchild)
        self.pre_order_recursion(root.rchild)

    def pre_order(self):
        s = Stack()
        p = self
        while not s.is_empty() or p is not None:
            if (p):
                print(p.data, end=' ')
                s.push(p)
                p = p.lchild
            else:
                p = s.pop()
                p = p.rchild

    def in_order_recursion(self, root):
        if root is None:
            return
        self.in_order_recursion(root.lchild)
        print(root.data, end=' ')
        self.in_order_recursion(root.rchild)

    def in_order(self):
        s = Stack()
        p = self
        while not s.is_empty() or p is not None:
            if (p):
                s.push(p)
                p = p.lchild
            else:
                p = s.pop()
                print(p.data, end=' ')
                p = p.rchild

    def post_order_recursion(self, root):
        if root is None:
            return
        self.post_order_recursion(root.lchild)
        self.post_order_recursion(root.rchild)
        print(root.data, end=' ')

    def post_order(self):
        s = Stack()
        p = self
        r = None
        while not s.is_empty() or p is not None:
            if (p):
                s.push(p)
                p = p.lchild
            else:
                p = s.stack[s.top]
                if p.rchild is not None and p.rchild is not r:
                    p = p.rchild
                    s.push(p)
                    p = p.lchild
                else:
                    p = s.pop()
                    print(p.data, end=' ')
                    r = p
                    p = None

    def get_height(self):
        if self.lchild is not None:
            hl = self.lchild.get_height()+1
        else:
            hl = 1

        if self.rchild is not None:
            hr = self.rchild.get_height()+1
        else:
            hr = 1
        h = hl if hl > hr else hr
        return h

if __name__ == "__main__":
    ll = Bi_tree(3)
    rr = Bi_tree(5)
    l = Bi_tree(2, lchhid=ll)
    r = Bi_tree(4, lchhid=None, rchid=rr)
    tree = Bi_tree(1, lchhid=l, rchid=r)
    tree.pre_order_recursion(tree)
    print()
    tree.pre_order()
    print()
    tree.post_order_recursion(tree)
    print()
    tree.post_order()
    print()
    print(tree.get_height())